# 自定义Token有效期
from rest_framework.authentication import TokenAuthentication
from rest_framework.exceptions import AuthenticationFailed
from rest_framework.authtoken.models import Token
import datetime
from rest_framework.response import Response

class ExpiringTokenAuthentication(TokenAuthentication):
    def authenticate_credentials(self, key):  # key是token
        try:
            token_obj = Token.objects.get(key=key)
            # Token过期时间为8h
            expire_date = token_obj.created + datetime.timedelta(hours=8)
            if  expire_date <= datetime.datetime.now():
                token_obj.delete()
                res = {'code': 501, 'msg': '会话过期，请重新登录'}
                raise AuthenticationFailed(res)
            return token_obj.user, token_obj.key
        except Token.DoesNotExist:
            res = {'code': 502, 'msg': 'Token不存在!'}
            raise AuthenticationFailed(res)